// 示例：简单的生产者-消费者向量加法
// 演示 Tawa 方言的 aref 操作与分区

module {
  // 向量加法函数
  func.func @vector_add(%a: tensor<16xf32>, %b: tensor<16xf32>) -> tensor<16xf32> {
    // 创建深度为 2 的 aref 缓冲区
    %aref = tawa.aref.create : !tawa.aref<f32, 2>
    
    // 生产者分区（WG0）：负责数据加载
    tawa.partition @producer partition<0> {
      // 将数据 a 放入 aref
      tawa.aref.put %aref, %a : !tawa.aref<f32, 2>, tensor<16xf32>
    }
    
    // 消费者分区（WG1）：负责计算
    tawa.partition @consumer partition<1> {
      // 从 aref 获取数据
      %data = tawa.aref.get %aref : !tawa.aref<f32, 2> -> tensor<16xf32>
      
      // 执行向量加法
      %result = arith.addf %data, %b : tensor<16xf32>
      
      // 通知数据已消费
      tawa.aref.consumed %aref : !tawa.aref<f32, 2>
      
      func.return %result : tensor<16xf32>
    }
  }
}
